# sense_voice_asr.py
import gc
from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess
import numpy as np
import torch
from config import app_config
from ..util.audio_util import load_audio

from .asr_base import BaseASR
from .asr_registry import register_asr



@register_asr("sense_voice_asr")
class SenseVoiceASR(BaseASR):
    def __init__(self):
        super().__init__()
        self.model_dir = app_config.SENSE_VOICE_ASR_MODEL
        self._model = None  # 延迟加载

    @property
    def model(self):
        """懒加载模型"""
        if self._model is None:
            print(f"[SenseVoiceASR] 正在加载模型: {self.model_dir}")
            self._model = AutoModel(
                model=self.model_dir,
                vad_model="fsmn-vad",
                vad_kwargs={"max_single_segment_time": 30000},
                device="cpu",  # 可改为配置项
            )
            print("[SenseVoiceASR] 模型加载完成")
        return self._model

    def transcribe(self, audio_data: np.ndarray = None, audio_path: str = None) -> str: # type: ignore
        if audio_data is None:
            if not audio_path:
                raise ValueError("audio_data is None, please provide audio_path")
            audio_data = load_audio(audio_path)

        res = self.model.generate(
            input=audio_data,
            cache={},
            language="auto",
            use_itn=True,
            batch_size_s=60,
            merge_vad=True,
            merge_length_s=15,
            ban_emo_unk=True,
        )
        text = rich_transcription_postprocess(res[0]["text"])
        return text

    def cleanup(self):
        """主动释放模型资源"""
        if self._model is not None:
            # FunASR 模型底层是 PyTorch，尝试删除并置空
            del self._model
            self._model = None
            print("[SenseVoiceASR] 模型已从内存中删除")

        # 清理 PyTorch 缓存（尤其是 GPU）
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            print("[SenseVoiceASR] CUDA 缓存已清空")

        # 触发 Python 垃圾回收
        gc.collect()
        print("[SenseVoiceASR] 资源清理完成")

    def __del__(self):
        """析构函数：对象被销毁时自动调用"""
        if self._model is not None:
            print(f"[SenseVoiceASR] 对象即将销毁，自动调用 cleanup...")
            self.cleanup()    